import torch
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from pymilvus import Collection

# -------------------------
# 1) 生成文本嵌入 (使用 GPT-2)
# -------------------------
def generate_embeddings(text: str, model_name: str = "gpt2") -> np.ndarray:
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)

    # 编码文本
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)

    # GPT-2 的嵌入是最后一层的 hidden state
    embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
    return embeddings

# -------------------------
# 2) 使用 Milvus 搜索并与 GPT-2 生成答案
# -------------------------
def rag_query(collection: Collection, query: str, top_k: int = 3, model_name: str = "gpt2") -> str:
    # 生成查询文本的 Embedding
    query_embedding = generate_embeddings(query, model_name=model_name)

    # 使用 Milvus 搜索相关文本
    search_params = {"metric_type": "COSINE", "params": {"nprobe": 16}}
    results = collection.search(data=[query_embedding], anns_field="embedding", param=search_params, limit=top_k, output_fields=["text"])

    # 获取检索到的上下文内容
    context = "\n".join([hit.entity["text"] for hit in results[0]])

    # 使用 GPT-2 生成最终答案
    input_text = f"问题：{query}\n\n相关内容：{context}\n\n回答："
    return generate_gpt2_answer(input_text, model_name)

# -------------------------
# 3) 使用 GPT-2 模型生成回答
# -------------------------
def generate_gpt2_answer(input_text: str, model_name: str = "gpt2") -> str:
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_name)

    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
    outputs = model.generate(inputs["input_ids"], max_length=500, num_return_sequences=1)

    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

query = "请输入你的问题"
top_k = 3  # 设置要检索的最相关文本数量
answer = rag_query(collection, query, top_k)
print("最终回答：", answer)